%	Compute the canonical example of Section IX of LANCIA, RUSSO, and WORRALL (2024) %
%   The value iteration is using ex-ante promised utility as state variable            %
%	Calls functions:    Utility.m                                                      %
%                       Autarky.m                                                      %
%                       Maxstat.m                                                      %
%                       Intconv.m                                                      %
%                       Policy_Iteration.m                                             %

clear all; close all;
tic
warning('off')

cd ..\
cd Matlab_functions

global wgrid betta delta gamma pi1 pi2 theta e1 e2 ey1 ey2 n v1aut v2aut wfb vfb nw w_st wmin wmax

%----------------------------------------------
% Parameter Setting
%----------------------------------------------

gamma = 1;                              % Risk Aversion
pi1 = 0.5;                              % prob state 1
pi2 = 1-pi1;                            % prob state 2
theta = 0;                              % aggregate shock  
e1 = 1+theta;                           % endowment state 1
e2 = 1-(pi1/(pi2))*theta;               % endowment state 2
sigma = 0.1;                            % measure of volatility of endowment
kappa = 3/5;                            % measure of mean of endowment
ey1 = (kappa-sigma*pi2/pi1)*e1;         % endowment state 1 young
eo1 = (1-kappa+sigma*pi2/pi1)*e1;       % endowment state 1 old
ey2 = (kappa+sigma)*e2;                 % endowment state 2 young
eo2 = (1-kappa-sigma)*e2;               % endowment state 2 old
betta = exp(-1/75);                     % individual discount Factor
delta = betta;                          % social planner discount Factor
n=0;                                    % fertility rate

%----------------------------------------------
% Checks
%----------------------------------------------

lamb1 = (eo1/ey1)^gamma;
lamb2 = (eo2/ey2)^gamma;
S1 = lamb1-betta;          % this must be positive, state 1 bad state for the young
S2 = lamb2-betta ;         % this must be negative, state 2 good state for the young
Sfb = lamb1-betta/delta;   % if negative, then FB implies positive transfers in both states
No_Samuelson = (1-kappa)-betta*kappa;   % if positive, then dynamic efficiency in deterministic case

%----------------------------------------------
% Find max and min of promised values
%----------------------------------------------

[v1aut,v2aut,waut] = Autarky(gamma,betta,pi1,pi2,ey1,ey2,eo1,eo2);
w_st = Maxstat(gamma,betta,pi1,pi2,ey1,ey2,eo1,eo2,v1aut,v2aut);

%----------------------------------------------
% Find the maximum interior convergence
%----------------------------------------------

[vfb,wfb,cy1fb,cy2fb,w_conv]=Intconv(gamma,betta,delta,n,pi1,pi2,v1aut);

%----------------------------------------------
% grid for promised values
%----------------------------------------------

wmin = waut;
wmax = w_st;
nw = 300; % search method for optimal grid number
wgrid = cheby_points(wmin,wmax,nw)';

%----------------------------------------------
% Iterating the Policy Function
%----------------------------------------------

distmia=100;
olddist=1;
iter=0;
tollv=10^(-6);
tic
fprintf('For this tolerance, I predict that about %g iterations will be needed \n',round(fsolve(@(x)tollv.^(1/x)-delta,3)));

cy1_old = cy1fb*ones(nw,1);
cy2_old = cy2fb*ones(nw,1);
w1_old = wfb*ones(nw,1);
w2_old = wfb*ones(nw,1);

% Policy iteration 
P_old=[cy1_old,cy2_old,w1_old,w2_old];
 
 while distmia > tollv
    iter=iter+1;
 
    [V,cy1,cy2,w1,w2,w_thr,w_approx] = Value_Iteration(cy1_old,cy2_old,w1_old,w2_old); 
    w_upper=fsolve(@(x)(x-interp1(wgrid,w2,x,'spline')),waut);

    P=[cy1,cy2,w1,w2]; 
    
    distmia=max(max(abs(P-P_old)));
    distmia=distmia+abs(cy1fb-interp1(wgrid,cy1,w_thr,'spline'));
    
    fprintf('Iteration %g completed, Distance is: %g   Improved by %9.1f %% \n',iter,distmia,(((olddist-distmia)/olddist))*100);
     
    olddist=distmia;
    P_old=P;
    cy1_old=cy1;
    cy2_old=cy2;
    w1_old=w1;
    w2_old=w2; 
end
    

figure(1)
plot(wgrid,w1,'r'); hold on;
plot(wgrid,w2,'b'); hold on;
plot(wgrid,wgrid,'black'); hold on;
plot([w_upper w_upper],[wmin wmax],'--','Color','k'); hold on;...
plot([w_thr w_thr],[wmin wmax],'--','Color','r'); hold off;...
title('Promised utilities','fontsize',08)

figure(2)                   
plot(wgrid,cy1./e1,'r'); hold on;
plot(wgrid,cy2./e2,'b'); xlim([w_thr w_upper]);
title('Young Consumption','fontsize',08)

figure(3)
plot(wgrid,V,'b','LineWidth',2); hold on;
line([w_thr w_thr],[V(end) max(V(abs(V(1)-V)<1e-10))],'LineStyle',':','LineWidth',0.5)
text(w_thr,V(end)-0.05,'$\omega_0$','interpreter','latex','FontSize',20)
ylim([V(end) V(1)+0.1]);

toc

% Ex-post promise
% next period consumption
cy11=interp1(wgrid,cy1,w1,'spline');
cy12=interp1(wgrid,cy2,w1,'spline');
cy21=interp1(wgrid,cy1,w2,'spline');
cy22=interp1(wgrid,cy2,w2,'spline');

omega1=utility(1-cy1,gamma);
omega2=utility(1-cy2,gamma);

omega11=utility(1-cy11,gamma);
omega12=utility(1-cy12,gamma);

omega21=utility(1-cy21,gamma);
omega22=utility(1-cy22,gamma);

cy1omega=interp1(wgrid,cy1,omega1,'spline');
cy2omega=interp1(wgrid,cy2,omega2,'spline');

figure(1)
subplot(2,2,1)
plot(omega1,omega11,'r',omega1,omega12,'b',omega1,omega1,'k:','LineWidth',0.5); hold on;
plot(log(1/2),log(1/2),'k*','LineWidth',0.5); hold on;
plot(omega1,log(1/2).*ones(1,length(omega1)),'k:','LineWidth',0.1); hold on;
hold off
legend('r=s(1)','r=s(2)')
set(gca,'FontSize',08);
title('Promised utility s_0=s(1)','FontSize',08)

subplot(2,2,2)
plot(omega2,omega21,'r',omega2,omega22,'b',omega2,omega2,'k:','LineWidth',0.5); hold on;
plot(log(1/2),log(1/2),'k*','LineWidth',0.5); hold on;
plot(omega2,log(1/2).*ones(1,length(omega2)),'k:','LineWidth',0.1); hold on;
hold off
legend('r=s(1)','r=s(2)')
set(gca,'FontSize',08);
title('Promised utility s_0=s(2)','FontSize',08)

%----------------------------------------------
% Outcome Variables
%----------------------------------------------

cd ..\
cd Stored_file

save('Canonical_policy.mat','V','wgrid', 'w_thr', 'V', 'w1', 'w2', 'cy1', 'cy2', 'delta', 'betta', 'gamma', 'theta', 'pi1', 'pi2', 'e1', 'e2', 'ey1', 'ey2', 'eo1', 'eo2')
